+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| Module                                                                                                                                                                           | Module params                    | Input                                                                                                                                                                                                                                                                                      | Output                                                                                                                                                                                                                                                      |   Param count |   Param bytes |
+==================================================================================================================================================================================+==================================+============================================================================================================================================================================================================================================================================================+=============================================================================================================================================================================================================================================================+===============+===============+
| alphafold (AlphaFold)                                                                                                                                                            |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | {'distogram': {'bin_edges': f32[63], 'logits': f32[631,631,64]},                                                                                                                                                                                            |    93,223,705 |     372.89 MB |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'experimentally_resolved': {'logits': f32[631,37]},                                                                                                                                                                                                        |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'masked_msa': {'logits': f32[252,631,22]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'predicted_aligned_error': {'asym_id': f32[631],                                                                                                                                                                                                           |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                              'breaks': f32[63],                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                              'logits': f32[631,631,64]},                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |  'predicted_lddt': {'logits': f32[631,50]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |  'structure_module': {'final_atom_mask': f32[631,37],                                                                                                                                                                                                       |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                       'final_atom_positions': f32[631,37,3]}}                                                                                                                                                                                               |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                               |                                  | batch={'aatype': s32[631],                                                                                                                                                                                                                                                                 | {'distogram': {'bin_edges': f32[63], 'logits': f32[631,631,64]},                                                                                                                                                                                            |    93,223,705 |     372.89 MB |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'experimentally_resolved': {'logits': f32[631,37]},                                                                                                                                                                                                        |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'masked_msa': {'logits': f32[252,631,22]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'predicted_aligned_error': {'asym_id': f32[631],                                                                                                                                                                                                           |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                              'breaks': f32[63],                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                              'logits': f32[631,631,64]},                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |  'predicted_lddt': {'logits': f32[631,50]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |  'representations': {'msa': f32[252,631,256],                                                                                                                                                                                                               |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                      'msa_first_row': f32[631,256],                                                                                                                                                                                                         |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                      'pair': f32[631,631,128],                                                                                                                                                                                                              |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                      'single': f32[631,384],                                                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                      'structure_module': f32[631,384]},                                                                                                                                                                                                     |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |  'structure_module': {'final_atom_mask': f32[631,37],                                                                                                                                                                                                       |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                       'final_atom_positions': f32[631,37,3]}}                                                                                                                                                                                               |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f2546d310>                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                                 |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | {'msa': f32[252,631,256],                                                                                                                                                                                                                                   |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'pair': f32[631,631,128],                                                                                                                                                                                                                                  |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'single': f32[631,384]}                                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_1d (Linear)                                                                                                                   | weights: f32[21,256]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,256]                                                                                                                                                                                                                                                |         5,632 |      22.53 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_msa (Linear)                                                                                                                  | weights: f32[49,256]             | f32[252,631,49]                                                                                                                                                                                                                                                                            | f32[252,631,256]                                                                                                                                                                                                                                            |        12,800 |      51.20 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/left_single (Linear)                                                                                                                     | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/right_single (Linear)                                                                                                                    | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pos_linear (Linear)                                                                                                                 | weights: f32[15,128]             | f32[631,631,15]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         2,048 |       8.19 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_msa_first_row_norm (LayerNorm)                                                                                                      | offset: f32[256]                 | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,256]                                                                                                                                                                                                                                                |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[256]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pair_norm (LayerNorm)                                                                                                               | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                              |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | f32[631,631,128]                                                                                                                                                                                                                                            |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[252,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_deletion_mean': f32[252,631],                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_profile': f32[252,631,23],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[252,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_bert_mask': f32[5891,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_deletion_matrix': f32[5891,631],                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa': s32[5891,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa_mask': f32[5891,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': f32[252,631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[252,631],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_profile': f32[631,22],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3],                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'true_msa': s32[252,631]}                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/~_relative_encoding/position_activations (Linear)                                                                                        | weights: f32[73,128]             | f32[631,631,73]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         9,472 |      37.89 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                   |                                  | is_training=False, multichain_mask_2d=bool[631,631], padding_mask_2d=f32[631,631], query_embedding=f32[631,631,128], safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f4c580150>, template_batch={'template_aatype': s32[4,631],                                                      | f32[631,631,128]                                                                                                                                                                                                                                            |       241,344 |     965.38 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'template_all_atom_positions': f32[4,631,37,3]}                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                   |                                  | f32[631,631,128], s32[631], f32[631,37,3], f32[631,37], f32[631,631], bool[631,631], False, <alphafold.model.prng.SafeKey object at 0x7f2f471b1d10>                                                                                                                                        | f32[631,631,64]                                                                                                                                                                                                                                             |       233,024 |     932.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/query_embedding_norm (LayerNorm)                                                            | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_0 (Linear)                                                          | weights: f32[39,64]              | f32[631,631,39]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         2,560 |      10.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_1 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_2 (Linear)                                                          | weights: f32[22,64]              | f32[1,631,22]                                                                                                                                                                                                                                                                              | f32[1,631,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_3 (Linear)                                                          | weights: f32[22,64]              | f32[631,1,22]                                                                                                                                                                                                                                                                              | f32[631,1,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_4 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_5 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_6 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_7 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_8 (Linear)                                                          | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                                 |                                  | ((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f2f4c68ea90>),), None                                                                                                                                                                                                        | (((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f2f4c6aba10>),),                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                   |                                  | act=f32[631,631,64], is_training=False, pair_mask=f32[631,631], safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f47de4450>                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)           | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/query_norm (LayerNorm)        | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)             | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/query_norm (LayerNorm)          | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                   |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        33,408 |     133.63 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/input_layer_norm (LayerNorm)                   | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/output_layer_norm (LayerNorm)                                                               | offset: f32[64]                  | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[64]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/output_linear (Linear)                                                                                                | weights: f32[64,128]             | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_activations (Linear)                                                                                                           | weights: f32[25,64]              | f32[1152,631,25]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |         1,664 |       6.66 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                              |                                  | (({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f2f45310590>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f324f556a50>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                     |                                  | activations={'msa': f32[1152,631,64], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[1152,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f37587f50>                                                                                   | {'msa': f32[1152,631,64], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                    | output_w: f32[4,32,32,128]       | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |       541,952 |       2.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_b: f32[4,128]             | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/layer_norm_input (LayerNorm)                                                                          | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/left_projection (Linear)                                                                              | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/right_projection (Linear)                                                                             | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                           | feat_2d_weights: f32[4,128,8]    | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        88,064 |     352.26 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                                  | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                                | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        82,432 |     329.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                   |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        54,272 |     217.09 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/query_norm (LayerNorm)                                                                       | offset: f32[4,64]                | f32[631,1152,64]                                                                                                                                                                                                                                                                           | f32[631,1152,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                              |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |       132,864 |     531.46 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/input_layer_norm (LayerNorm)                                                                              | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                     | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/query_norm (LayerNorm)                                                                  | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                       | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/query_norm (LayerNorm)                                                                    | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                             |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       527,872 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/input_layer_norm (LayerNorm)                                                                             | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_single_embedding (Linear)                                                                                                       | weights: f32[34,256]             | f32[4,631,34]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |         8,960 |      35.84 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_projection (Linear)                                                                                                             | weights: f32[256,256]            | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                            |                                  | (({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f2f360d2310>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f2f360d2490>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                                 |                                  | activations={'msa': f32[256,631,256], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[256,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f35a2db10>                                                                                    | {'msa': f32[256,631,256], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                                | output_w: f32[48,32,32,128]      | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     7,111,680 |      28.45 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_b: f32[48,128]            | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/layer_norm_input (LayerNorm)                                                                      | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/left_projection (Linear)                                                                          | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/right_projection (Linear)                                                                         | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                       | feat_2d_weights: f32[48,128,8]   | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,839,232 |      63.36 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                              | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                            | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                            |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,777,792 |      63.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/query_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[631,256,256]                                                                                                                                                                                                                                                                           | f32[631,256,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                          |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    25,251,840 |     101.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/input_layer_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                                 | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/query_norm (LayerNorm)                                                              | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                   | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/query_norm (LayerNorm)                                                                | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                         |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     6,334,464 |      25.34 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/input_layer_norm (LayerNorm)                                                                         | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/single_activations (Linear)                                                                                                              | weights: f32[256,384]            | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |        98,688 |     394.75 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                                 |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | {'msa': f32[252,631,256],                                                                                                                                                                                                                                   |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'pair': f32[631,631,128],                                                                                                                                                                                                                                  |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'single': f32[631,384]}                                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  | safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f272d40d0>                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_1d (Linear)                                                                                                                   | weights: f32[21,256]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,256]                                                                                                                                                                                                                                                |         5,632 |      22.53 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_msa (Linear)                                                                                                                  | weights: f32[49,256]             | f32[252,631,49]                                                                                                                                                                                                                                                                            | f32[252,631,256]                                                                                                                                                                                                                                            |        12,800 |      51.20 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/left_single (Linear)                                                                                                                     | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/right_single (Linear)                                                                                                                    | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pos_linear (Linear)                                                                                                                 | weights: f32[15,128]             | f32[631,631,15]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         2,048 |       8.19 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_msa_first_row_norm (LayerNorm)                                                                                                      | offset: f32[256]                 | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,256]                                                                                                                                                                                                                                                |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[256]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pair_norm (LayerNorm)                                                                                                               | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                              |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | f32[631,631,128]                                                                                                                                                                                                                                            |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[252,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_deletion_mean': f32[252,631],                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_profile': f32[252,631,23],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[252,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_bert_mask': f32[5891,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_deletion_matrix': f32[5891,631],                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa': s32[5891,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa_mask': f32[5891,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': f32[252,631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[252,631],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_profile': f32[631,22],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3],                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'true_msa': s32[252,631]}                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/~_relative_encoding/position_activations (Linear)                                                                                        | weights: f32[73,128]             | f32[631,631,73]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         9,472 |      37.89 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                   |                                  | is_training=False, multichain_mask_2d=bool[631,631], padding_mask_2d=f32[631,631], query_embedding=f32[631,631,128], safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f34096dd0>, template_batch={'template_aatype': s32[4,631],                                                      | f32[631,631,128]                                                                                                                                                                                                                                            |       241,344 |     965.38 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'template_all_atom_positions': f32[4,631,37,3]}                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                   |                                  | f32[631,631,128], s32[631], f32[631,37,3], f32[631,37], f32[631,631], bool[631,631], False, <alphafold.model.prng.SafeKey object at 0x7f2f34173dd0>                                                                                                                                        | f32[631,631,64]                                                                                                                                                                                                                                             |       233,024 |     932.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/query_embedding_norm (LayerNorm)                                                            | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_0 (Linear)                                                          | weights: f32[39,64]              | f32[631,631,39]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         2,560 |      10.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_1 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_2 (Linear)                                                          | weights: f32[22,64]              | f32[1,631,22]                                                                                                                                                                                                                                                                              | f32[1,631,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_3 (Linear)                                                          | weights: f32[22,64]              | f32[631,1,22]                                                                                                                                                                                                                                                                              | f32[631,1,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_4 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_5 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_6 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_7 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_8 (Linear)                                                          | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                                 |                                  | ((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f2f35d64110>),), None                                                                                                                                                                                                        | (((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f2f35d677d0>),),                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                   |                                  | act=f32[631,631,64], is_training=False, pair_mask=f32[631,631], safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f344bc910>                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)           | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/query_norm (LayerNorm)        | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)             | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/query_norm (LayerNorm)          | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                   |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        33,408 |     133.63 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/input_layer_norm (LayerNorm)                   | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/output_layer_norm (LayerNorm)                                                               | offset: f32[64]                  | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[64]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/output_linear (Linear)                                                                                                | weights: f32[64,128]             | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_activations (Linear)                                                                                                           | weights: f32[25,64]              | f32[1152,631,25]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |         1,664 |       6.66 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                              |                                  | (({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f2f2cd3ac10>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f2f2d117d10>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                     |                                  | activations={'msa': f32[1152,631,64], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[1152,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f2c54cdd0>                                                                                   | {'msa': f32[1152,631,64], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                    | output_w: f32[4,32,32,128]       | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |       541,952 |       2.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_b: f32[4,128]             | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/layer_norm_input (LayerNorm)                                                                          | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/left_projection (Linear)                                                                              | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/right_projection (Linear)                                                                             | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                           | feat_2d_weights: f32[4,128,8]    | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        88,064 |     352.26 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                                  | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                                | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        82,432 |     329.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                   |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        54,272 |     217.09 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/query_norm (LayerNorm)                                                                       | offset: f32[4,64]                | f32[631,1152,64]                                                                                                                                                                                                                                                                           | f32[631,1152,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                              |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |       132,864 |     531.46 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/input_layer_norm (LayerNorm)                                                                              | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                     | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/query_norm (LayerNorm)                                                                  | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                       | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/query_norm (LayerNorm)                                                                    | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                             |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       527,872 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/input_layer_norm (LayerNorm)                                                                             | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_single_embedding (Linear)                                                                                                       | weights: f32[34,256]             | f32[4,631,34]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |         8,960 |      35.84 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_projection (Linear)                                                                                                             | weights: f32[256,256]            | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                            |                                  | (({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f2f268824d0>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f2f26882910>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                                 |                                  | activations={'msa': f32[256,631,256], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[256,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f25ec4d10>                                                                                    | {'msa': f32[256,631,256], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                                | output_w: f32[48,32,32,128]      | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     7,111,680 |      28.45 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_b: f32[48,128]            | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/layer_norm_input (LayerNorm)                                                                      | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/left_projection (Linear)                                                                          | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/right_projection (Linear)                                                                         | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                       | feat_2d_weights: f32[48,128,8]   | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,839,232 |      63.36 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                              | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                            | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                            |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,777,792 |      63.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/query_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[631,256,256]                                                                                                                                                                                                                                                                           | f32[631,256,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                          |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    25,251,840 |     101.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/input_layer_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                                 | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/query_norm (LayerNorm)                                                              | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                   | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/query_norm (LayerNorm)                                                                | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                         |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     6,334,464 |      25.34 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/input_layer_norm (LayerNorm)                                                                         | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/single_activations (Linear)                                                                                                              | weights: f32[256,384]            | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |        98,688 |     394.75 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                                 |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'act': f32[631,384],                                                                                                                                                                                                                                       |     2,018,540 |       8.07 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |  'final_atom_mask': f32[631,37],                                                                                                                                                                                                                            |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |  'final_atom_positions': f32[631,37,3]}                                                                                                                                                                                                                     |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/single_layer_norm (LayerNorm)                                                                                                     | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/initial_projection (Linear)                                                                                                       | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/pair_layer_norm (LayerNorm)                                                                                                       | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f44cff190>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f25b6af50>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f2556b1d0>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f24f69290>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2f25a74610>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f332939be90>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f240595b410>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f2405231c50>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/distogram_head (DistogramHead)                                                                                                                     |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'bin_edges': f32[63], 'logits': f32[631,631,64]}                                                                                                                                                                                                           |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/distogram_head/half_logits (Linear)                                                                                                                | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/distogram_head (DistogramHead)                                                                                                                  | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/masked_msa_head (MaskedMsaHead)                                                                                                                    |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[252,631,22]}                                                                                                                                                                                                                                 |         5,654 |      22.62 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/masked_msa_head/logits (Linear)                                                                                                                    | weights: f32[256,22]             | f32[252,631,256]                                                                                                                                                                                                                                                                           | f32[252,631,22]                                                                                                                                                                                                                                             |         5,654 |      22.62 KB |
|  └ alphafold/alphafold_iteration/masked_msa_head (MaskedMsaHead)                                                                                                                 | bias: f32[22]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                            |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[631,50]}                                                                                                                                                                                                                                     |        73,010 |     292.04 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/input_layer_norm (LayerNorm)                                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/act_0 (Linear)                                                                                                                 | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/act_1 (Linear)                                                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/logits (Linear)                                                                                                                | weights: f32[128,50]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,50]                                                                                                                                                                                                                                                 |         6,450 |      25.80 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[50]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/experimentally_resolved_head (ExperimentallyResolvedHead)                                                                                          |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[631,37]}                                                                                                                                                                                                                                     |        14,245 |      56.98 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/experimentally_resolved_head/logits (Linear)                                                                                                       | weights: f32[384,37]             | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,37]                                                                                                                                                                                                                                                 |        14,245 |      56.98 KB |
|  └ alphafold/alphafold_iteration/experimentally_resolved_head (ExperimentallyResolvedHead)                                                                                       | bias: f32[37]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_aligned_error_head (PredictedAlignedErrorHead)                                                                                           |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'breaks': f32[63], 'logits': f32[631,631,64]}                                                                                                                                                                                                              |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_aligned_error_head/logits (Linear)                                                                                                       | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/predicted_aligned_error_head (PredictedAlignedErrorHead)                                                                                        | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                               |                                  | batch={'aatype': s32[631],                                                                                                                                                                                                                                                                 | {'distogram': {'bin_edges': f32[63], 'logits': f32[631,631,64]},                                                                                                                                                                                            |    93,223,705 |     372.89 MB |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'experimentally_resolved': {'logits': f32[631,37]},                                                                                                                                                                                                        |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'masked_msa': {'logits': f32[252,631,22]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'predicted_aligned_error': {'asym_id': f32[631],                                                                                                                                                                                                           |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                              'breaks': f32[63],                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                              'logits': f32[631,631,64]},                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |  'predicted_lddt': {'logits': f32[631,50]},                                                                                                                                                                                                                 |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |  'representations': {'msa': f32[252,631,256],                                                                                                                                                                                                               |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                      'msa_first_row': f32[631,256],                                                                                                                                                                                                         |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                      'pair': f32[631,631,128],                                                                                                                                                                                                              |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                      'single': f32[631,384],                                                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                      'structure_module': f32[631,384]},                                                                                                                                                                                                     |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |  'structure_module': {'final_atom_mask': f32[631,37],                                                                                                                                                                                                       |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                       'final_atom_positions': f32[631,37,3]}}                                                                                                                                                                                               |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23efdfc3d0>                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                                 |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | {'msa': f32[252,631,256],                                                                                                                                                                                                                                   |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'pair': f32[631,631,128],                                                                                                                                                                                                                                  |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'single': f32[631,384]}                                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_1d (Linear)                                                                                                                   | weights: f32[21,256]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,256]                                                                                                                                                                                                                                                |         5,632 |      22.53 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_msa (Linear)                                                                                                                  | weights: f32[49,256]             | f32[252,631,49]                                                                                                                                                                                                                                                                            | f32[252,631,256]                                                                                                                                                                                                                                            |        12,800 |      51.20 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/left_single (Linear)                                                                                                                     | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/right_single (Linear)                                                                                                                    | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pos_linear (Linear)                                                                                                                 | weights: f32[15,128]             | f32[631,631,15]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         2,048 |       8.19 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_msa_first_row_norm (LayerNorm)                                                                                                      | offset: f32[256]                 | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,256]                                                                                                                                                                                                                                                |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[256]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pair_norm (LayerNorm)                                                                                                               | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                              |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | f32[631,631,128]                                                                                                                                                                                                                                            |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[252,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_deletion_mean': f32[252,631],                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_profile': f32[252,631,23],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[252,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_bert_mask': f32[5891,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_deletion_matrix': f32[5891,631],                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa': s32[5891,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa_mask': f32[5891,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': f32[252,631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[252,631],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_profile': f32[631,22],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3],                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'true_msa': s32[252,631]}                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/~_relative_encoding/position_activations (Linear)                                                                                        | weights: f32[73,128]             | f32[631,631,73]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         9,472 |      37.89 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                   |                                  | is_training=False, multichain_mask_2d=bool[631,631], padding_mask_2d=f32[631,631], query_embedding=f32[631,631,128], safe_key=<alphafold.model.prng.SafeKey object at 0x7f23ff97fcd0>, template_batch={'template_aatype': s32[4,631],                                                      | f32[631,631,128]                                                                                                                                                                                                                                            |       241,344 |     965.38 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'template_all_atom_positions': f32[4,631,37,3]}                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                   |                                  | f32[631,631,128], s32[631], f32[631,37,3], f32[631,37], f32[631,631], bool[631,631], False, <alphafold.model.prng.SafeKey object at 0x7f23ffa2c9d0>                                                                                                                                        | f32[631,631,64]                                                                                                                                                                                                                                             |       233,024 |     932.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/query_embedding_norm (LayerNorm)                                                            | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_0 (Linear)                                                          | weights: f32[39,64]              | f32[631,631,39]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         2,560 |      10.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_1 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_2 (Linear)                                                          | weights: f32[22,64]              | f32[1,631,22]                                                                                                                                                                                                                                                                              | f32[1,631,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_3 (Linear)                                                          | weights: f32[22,64]              | f32[631,1,22]                                                                                                                                                                                                                                                                              | f32[631,1,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_4 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_5 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_6 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_7 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_8 (Linear)                                                          | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                                 |                                  | ((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f24047baa50>),), None                                                                                                                                                                                                        | (((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f24045fa090>),),                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                   |                                  | act=f32[631,631,64], is_training=False, pair_mask=f32[631,631], safe_key=<alphafold.model.prng.SafeKey object at 0x7f32a0ab49d0>                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)           | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/query_norm (LayerNorm)        | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)             | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/query_norm (LayerNorm)          | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                   |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        33,408 |     133.63 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/input_layer_norm (LayerNorm)                   | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/output_layer_norm (LayerNorm)                                                               | offset: f32[64]                  | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[64]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/output_linear (Linear)                                                                                                | weights: f32[64,128]             | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_activations (Linear)                                                                                                           | weights: f32[25,64]              | f32[1152,631,25]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |         1,664 |       6.66 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                              |                                  | (({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f23fae68a50>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f23fae68c90>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                     |                                  | activations={'msa': f32[1152,631,64], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[1152,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23fa68d550>                                                                                   | {'msa': f32[1152,631,64], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                    | output_w: f32[4,32,32,128]       | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |       541,952 |       2.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_b: f32[4,128]             | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/layer_norm_input (LayerNorm)                                                                          | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/left_projection (Linear)                                                                              | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/right_projection (Linear)                                                                             | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                           | feat_2d_weights: f32[4,128,8]    | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        88,064 |     352.26 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                                  | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                                | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        82,432 |     329.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                   |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        54,272 |     217.09 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/query_norm (LayerNorm)                                                                       | offset: f32[4,64]                | f32[631,1152,64]                                                                                                                                                                                                                                                                           | f32[631,1152,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                              |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |       132,864 |     531.46 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/input_layer_norm (LayerNorm)                                                                              | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                     | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/query_norm (LayerNorm)                                                                  | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                       | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/query_norm (LayerNorm)                                                                    | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                             |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       527,872 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/input_layer_norm (LayerNorm)                                                                             | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_single_embedding (Linear)                                                                                                       | weights: f32[34,256]             | f32[4,631,34]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |         8,960 |      35.84 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_projection (Linear)                                                                                                             | weights: f32[256,256]            | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                            |                                  | (({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f23f88dbf90>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f2f4512c490>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                                 |                                  | activations={'msa': f32[256,631,256], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[256,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f8170690>                                                                                    | {'msa': f32[256,631,256], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                                | output_w: f32[48,32,32,128]      | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     7,111,680 |      28.45 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_b: f32[48,128]            | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/layer_norm_input (LayerNorm)                                                                      | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/left_projection (Linear)                                                                          | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/right_projection (Linear)                                                                         | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                       | feat_2d_weights: f32[48,128,8]   | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,839,232 |      63.36 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                              | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                            | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                            |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,777,792 |      63.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/query_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[631,256,256]                                                                                                                                                                                                                                                                           | f32[631,256,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                          |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    25,251,840 |     101.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/input_layer_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                                 | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/query_norm (LayerNorm)                                                              | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                   | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/query_norm (LayerNorm)                                                                | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                         |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     6,334,464 |      25.34 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/input_layer_norm (LayerNorm)                                                                         | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/single_activations (Linear)                                                                                                              | weights: f32[256,384]            | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |        98,688 |     394.75 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                                 |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | {'msa': f32[252,631,256],                                                                                                                                                                                                                                   |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |  'pair': f32[631,631,128],                                                                                                                                                                                                                                  |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |  'single': f32[631,384]}                                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  | safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f24c0610>                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_1d (Linear)                                                                                                                   | weights: f32[21,256]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,256]                                                                                                                                                                                                                                                |         5,632 |      22.53 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/preprocess_msa (Linear)                                                                                                                  | weights: f32[49,256]             | f32[252,631,49]                                                                                                                                                                                                                                                                            | f32[252,631,256]                                                                                                                                                                                                                                            |        12,800 |      51.20 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/left_single (Linear)                                                                                                                     | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/right_single (Linear)                                                                                                                    | weights: f32[21,128]             | f32[631,21]                                                                                                                                                                                                                                                                                | f32[631,128]                                                                                                                                                                                                                                                |         2,816 |      11.26 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pos_linear (Linear)                                                                                                                 | weights: f32[15,128]             | f32[631,631,15]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         2,048 |       8.19 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_msa_first_row_norm (LayerNorm)                                                                                                      | offset: f32[256]                 | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,256]                                                                                                                                                                                                                                                |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[256]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/prev_pair_norm (LayerNorm)                                                                                                               | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                              |                                  | {'aatype': s32[631],                                                                                                                                                                                                                                                                       | f32[631,631,128]                                                                                                                                                                                                                                            |    91,095,744 |     364.38 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[252,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_deletion_mean': f32[252,631],                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_profile': f32[252,631,23],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[252,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_bert_mask': f32[5891,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_deletion_matrix': f32[5891,631],                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa': s32[5891,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'extra_msa_mask': f32[5891,631],                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': f32[252,631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[252,631],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_profile': f32[631,22],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3],                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'true_msa': s32[252,631]}                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/~_relative_encoding/position_activations (Linear)                                                                                        | weights: f32[73,128]             | f32[631,631,73]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         9,472 |      37.89 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer._relative_encoding)                                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                   |                                  | is_training=False, multichain_mask_2d=bool[631,631], padding_mask_2d=f32[631,631], query_embedding=f32[631,631,128], safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f77bead0>, template_batch={'template_aatype': s32[4,631],                                                      | f32[631,631,128]                                                                                                                                                                                                                                            |       241,344 |     965.38 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'template_all_atom_positions': f32[4,631,37,3]}                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                   |                                  | f32[631,631,128], s32[631], f32[631,37,3], f32[631,37], f32[631,631], bool[631,631], False, <alphafold.model.prng.SafeKey object at 0x7f23f782ce90>                                                                                                                                        | f32[631,631,64]                                                                                                                                                                                                                                             |       233,024 |     932.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/query_embedding_norm (LayerNorm)                                                            | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_0 (Linear)                                                          | weights: f32[39,64]              | f32[631,631,39]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         2,560 |      10.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_1 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_2 (Linear)                                                          | weights: f32[22,64]              | f32[1,631,22]                                                                                                                                                                                                                                                                              | f32[1,631,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_3 (Linear)                                                          | weights: f32[22,64]              | f32[631,1,22]                                                                                                                                                                                                                                                                              | f32[631,1,64]                                                                                                                                                                                                                                               |         1,472 |       5.89 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_4 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_5 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_6 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_7 (Linear)                                                          | bias: f32[64]                    | f32[631,631]                                                                                                                                                                                                                                                                               | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | weights: f32[64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_pair_embedding_8 (Linear)                                                          | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                                 |                                  | ((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f23fe6c6d50>),), None                                                                                                                                                                                                        | (((f32[631,631,64], <alphafold.model.prng.SafeKey object at 0x7f23fcfce710>),),                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                   |                                  | act=f32[631,631,64], is_training=False, pair_mask=f32[631,631], safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f7a79d10>                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |       218,240 |     872.96 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_outgoing (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)      |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        50,432 |     201.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)  | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_projection (Linear)      | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_projection (Linear)     | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/left_gate (Linear)            | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/right_gate (Linear)           | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm) | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/output_projection (Linear)    | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming/gating_linear (Linear)        | weights: f32[2,64,64]            | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_multiplication_incoming (TriangleMultiplication)   | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)           | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/query_norm (LayerNorm)        | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node/attention (Attention)         | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_starting_node (TriangleAttention)        | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)             | feat_2d_weights: f32[2,64,4]     | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        41,984 |     167.94 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/query_norm (LayerNorm)          | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,216 |     164.86 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node/attention (Attention)           | gating_w: f32[2,64,4,16]         | f32[3,631,64], f32[3,631,64], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                               | f32[3,631,64]                                                                                                                                                                                                                                               |        41,728 |     166.91 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/triangle_attention_ending_node (TriangleAttention)          | key_w: f32[2,64,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                | output_w: f32[2,4,16,64]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              | query_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | value_w: f32[2,64,4,16]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | gating_b: f32[2,4,16]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | output_b: f32[2,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                   |                                  | f32[631,631,64], f32[631,631]                                                                                                                                                                                                                                                              | f32[631,631,64]                                                                                                                                                                                                                                             |        33,408 |     133.63 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/input_layer_norm (LayerNorm)                   | offset: f32[2,64]                | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                | scale: f32[2,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                        |                                  | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,64]                                                                                                                                                                                                                                               |        33,152 |     132.61 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition1 (Linear)                           | weights: f32[2,64,128]           | f32[3,631,64]                                                                                                                                                                                                                                                                              | f32[3,631,128]                                                                                                                                                                                                                                              |        16,640 |      66.56 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/transition2 (Linear)                           | weights: f32[2,128,64]           | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,64]                                                                                                                                                                                                                                               |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition/sequential (Sequential)                     | bias: f32[2,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration/pair_transition (Transition)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/template_embedding_iteration (TemplateEmbeddingIteration)                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/__layer_stack_no_state (_LayerStackNoState)                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding/output_layer_norm (LayerNorm)                                                               | offset: f32[64]                  | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,64]                                                                                                                                                                                                                                             |           128 |      512.00 B |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding/single_template_embedding (SingleTemplateEmbedding)                                                                | scale: f32[64]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_embedding/output_linear (Linear)                                                                                                | weights: f32[64,128]             | f32[631,631,64]                                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/template_embedding (TemplateEmbedding)                                                                                                | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_activations (Linear)                                                                                                           | weights: f32[25,64]              | f32[1152,631,25]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |         1,664 |       6.66 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                              |                                  | (({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[1152,631,64], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f23f40c37d0>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f23f40c3510>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                     |                                  | activations={'msa': f32[1152,631,64], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[1152,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f37019d0>                                                                                   | {'msa': f32[1152,631,64], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |     2,805,248 |      11.22 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                    | output_w: f32[4,32,32,128]       | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[631,631,128]                                                                                                                                                                                                                                            |       541,952 |       2.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_b: f32[4,128]             | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/layer_norm_input (LayerNorm)                                                                          | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/left_projection (Linear)                                                                              | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean/right_projection (Linear)                                                                             | weights: f32[4,64,32]            | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,32]                                                                                                                                                                                                                                            |         8,320 |      33.28 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/outer_product_mean (OuterProductMean)                                                                                 | bias: f32[4,32]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                           | feat_2d_weights: f32[4,128,8]    | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        88,064 |     352.26 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                                  | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                                | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        82,432 |     329.73 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias/attention (Attention)                                                                   | gating_w: f32[4,64,8,8]          | f32[4,631,64], f32[4,631,64], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                               | f32[4,631,64]                                                                                                                                                                                                                                               |        86,528 |     346.11 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                        | key_w: f32[4,64,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                   |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |        54,272 |     217.09 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/query_norm (LayerNorm)                                                                       | offset: f32[4,64]                | f32[631,1152,64]                                                                                                                                                                                                                                                                           | f32[631,1152,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[4,1152,64], f32[4,1152,64], f32[4,1152,1], f32[4,1,1,1152]                                                                                                                                                                                                                             | f32[4,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention/attention (GlobalAttention)                                                                  | gating_w: f32[4,64,8,8]          | f32[3,1152,64], f32[3,1152,64], f32[3,1152,1], f32[3,1,1,1152]                                                                                                                                                                                                                             | f32[3,1152,64]                                                                                                                                                                                                                                              |        53,760 |     215.04 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_column_global_attention (MSAColumnGlobalAttention)                                                                | output_w: f32[4,8,8,64]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | query_w: f32[4,64,8,8]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | key_w: f32[4,64,8]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,64,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,8,8]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,64]              |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                              |                                  | f32[1152,631,64], f32[1152,631]                                                                                                                                                                                                                                                            | f32[1152,631,64]                                                                                                                                                                                                                                            |       132,864 |     531.46 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/input_layer_norm (LayerNorm)                                                                              | offset: f32[4,64]                | f32[1152,631,64]                                                                                                                                                                                                                                                                           | f32[1152,631,64]                                                                                                                                                                                                                                            |           512 |       2.05 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           | scale: f32[4,64]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                   |                                  | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,64]                                                                                                                                                                                                                                               |       132,352 |     529.41 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition1 (Linear)                                                                                      | weights: f32[4,64,256]           | f32[4,631,64]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |        66,560 |     266.24 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,256]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/transition2 (Linear)                                                                                      | weights: f32[4,256,64]           | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,64]                                                                                                                                                                                                                                               |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition/sequential (Sequential)                                                                                | bias: f32[4,64]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/msa_transition (Transition)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_outgoing (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                                |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       398,336 |       1.59 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                            | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_projection (Linear)                                                                | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_projection (Linear)                                                               | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/left_gate (Linear)                                                                      | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/right_gate (Linear)                                                                     | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                           | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/output_projection (Linear)                                                              | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming/gating_linear (Linear)                                                                  | weights: f32[4,128,128]          | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        66,048 |     264.19 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_multiplication_incoming (TriangleMultiplication)                                                             | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                     | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/query_norm (LayerNorm)                                                                  | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node/attention (Attention)                                                                   | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_starting_node (TriangleAttention)                                                                  | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                       | feat_2d_weights: f32[4,128,4]    | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       331,776 |       1.33 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/query_norm (LayerNorm)                                                                    | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       328,704 |       1.31 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node/attention (Attention)                                                                     | gating_w: f32[4,128,4,32]        | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       330,752 |       1.32 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/triangle_attention_ending_node (TriangleAttention)                                                                    | key_w: f32[4,128,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  | output_w: f32[4,4,32,128]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           | query_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[4,128,4,32]         |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[4,4,32]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[4,128]             |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                             |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |       527,872 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/input_layer_norm (LayerNorm)                                                                             | offset: f32[4,128]               | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |         1,024 |       4.10 KB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          | scale: f32[4,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                                  |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       526,848 |       2.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition1 (Linear)                                                                                     | weights: f32[4,128,512]          | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |       264,192 |       1.06 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,512]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/transition2 (Linear)                                                                                     | weights: f32[4,512,128]          | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |       262,656 |       1.05 MB |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition/sequential (Sequential)                                                                               | bias: f32[4,128]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack/pair_transition (Transition)                                                                                          |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/extra_msa_stack (EvoformerIteration)                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state (_LayerStackNoState)                                                                                           |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_single_embedding (Linear)                                                                                                       | weights: f32[34,256]             | f32[4,631,34]                                                                                                                                                                                                                                                                              | f32[4,631,256]                                                                                                                                                                                                                                              |         8,960 |      35.84 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/template_projection (Linear)                                                                                                             | weights: f32[256,256]            | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |        65,792 |     263.17 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[256]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                            |                                  | (({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                                                     | ((({'msa': f32[256,631,256], 'pair': f32[631,631,128]},                                                                                                                                                                                                     |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |   <alphafold.model.prng.SafeKey object at 0x7f23f1a7b1d0>),), None                                                                                                                                                                                                                         |    <alphafold.model.prng.SafeKey object at 0x7f23f1a7bb50>),),                                                                                                                                                                                              |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  None)                                                                                                                                                                                                                                                      |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                                 |                                  | activations={'msa': f32[256,631,256], 'pair': f32[631,631,128]}, is_training=False, masks={'msa': f32[256,631], 'pair': f32[631,631]}, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f1082990>                                                                                    | {'msa': f32[256,631,256], 'pair': f32[631,631,128]}                                                                                                                                                                                                         |    87,837,696 |     351.35 MB |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                                | output_w: f32[48,32,32,128]      | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     7,111,680 |      28.45 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_b: f32[48,128]            | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/layer_norm_input (LayerNorm)                                                                      | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/left_projection (Linear)                                                                          | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean/right_projection (Linear)                                                                         | weights: f32[48,256,32]          | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,32]                                                                                                                                                                                                                                             |       394,752 |       1.58 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/outer_product_mean (OuterProductMean)                                                                             | bias: f32[48,32]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                       | feat_2d_weights: f32[48,128,8]   | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,839,232 |      63.36 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False, pair_act=f32[631,631,128]                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/query_norm (LayerNorm)                                                              | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/feat_2d_norm (LayerNorm)                                                            | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias/attention (Attention)                                                               | gating_w: f32[48,256,8,32]       | f32[4,631,256], f32[4,631,256], f32[4,1,1,631], f32[8,631,631]                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    15,802,368 |      63.21 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_row_attention_with_pair_bias (MSARowAttentionWithPairBias)                                                    | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                            |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    15,777,792 |      63.11 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/query_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[631,256,256]                                                                                                                                                                                                                                                                           | f32[631,256,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[4,256,256], f32[4,256,256], f32[4,1,1,256]                                                                                                                                                                                                                                             | f32[4,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention/attention (Attention)                                                                           | gating_w: f32[48,256,8,32]       | f32[3,256,256], f32[3,256,256], f32[3,1,1,256]                                                                                                                                                                                                                                             | f32[3,256,256]                                                                                                                                                                                                                                              |    15,753,216 |      63.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_column_attention (MSAColumnAttention)                                                                         | key_w: f32[48,256,8,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,8,32,256]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,256,8,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,8,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,256]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                          |                                  | f32[256,631,256], f32[256,631]                                                                                                                                                                                                                                                             | f32[256,631,256]                                                                                                                                                                                                                                            |    25,251,840 |     101.01 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/input_layer_norm (LayerNorm)                                                                          | offset: f32[48,256]              | f32[256,631,256]                                                                                                                                                                                                                                                                           | f32[256,631,256]                                                                                                                                                                                                                                            |        24,576 |      98.30 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       | scale: f32[48,256]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                               |                                  | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,256]                                                                                                                                                                                                                                              |    25,227,264 |     100.91 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition1 (Linear)                                                                                  | weights: f32[48,256,1024]        | f32[4,631,256]                                                                                                                                                                                                                                                                             | f32[4,631,1024]                                                                                                                                                                                                                                             |    12,632,064 |      50.53 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,1024]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/transition2 (Linear)                                                                                  | weights: f32[48,1024,256]        | f32[4,631,1024]                                                                                                                                                                                                                                                                            | f32[4,631,256]                                                                                                                                                                                                                                              |    12,595,200 |      50.38 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition/sequential (Sequential)                                                                            | bias: f32[48,256]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/msa_transition (Transition)                                                                                       |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_outgoing (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                            |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     4,780,032 |      19.12 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/layer_norm_input (LayerNorm)                                                        | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_projection (Linear)                                                            | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_projection (Linear)                                                           | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/left_gate (Linear)                                                                  | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/right_gate (Linear)                                                                 | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/center_layer_norm (LayerNorm)                                                       | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/output_projection (Linear)                                                          | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming/gating_linear (Linear)                                                              | weights: f32[48,128,128]         | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |       792,576 |       3.17 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_multiplication_incoming (TriangleMultiplication)                                                         | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                                 | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/query_norm (LayerNorm)                                                              | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node/attention (Attention)                                                               | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_starting_node (TriangleAttention)                                                              | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                   | feat_2d_weights: f32[48,128,4]   | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     3,981,312 |      15.93 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/query_norm (LayerNorm)                                                                | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,944,448 |      15.78 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[4,631,128], f32[4,631,128], f32[4,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node/attention (Attention)                                                                 | gating_w: f32[48,128,4,32]       | f32[3,631,128], f32[3,631,128], f32[3,1,1,631], f32[4,631,631]                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,969,024 |      15.88 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/triangle_attention_ending_node (TriangleAttention)                                                                | key_w: f32[48,128,4,32]          |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              | output_w: f32[48,4,32,128]       |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         | query_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | value_w: f32[48,128,4,32]        |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            | gating_b: f32[48,4,32]           |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         | output_b: f32[48,128]            |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                         |                                  | f32[631,631,128], f32[631,631]                                                                                                                                                                                                                                                             | f32[631,631,128]                                                                                                                                                                                                                                            |     6,334,464 |      25.34 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  | is_training=False                                                                                                                                                                                                                                                                          |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/input_layer_norm (LayerNorm)                                                                         | offset: f32[48,128]              | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |        12,288 |      49.15 KB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      | scale: f32[48,128]               |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[4,631,128]                                                                                                                                                                                                                                                                             | f32[4,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[4,631,512]                                                                                                                                                                                                                                                                             | f32[4,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                              |                                  | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     6,322,176 |      25.29 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition1 (Linear)                                                                                 | weights: f32[48,128,512]         | f32[3,631,128]                                                                                                                                                                                                                                                                             | f32[3,631,512]                                                                                                                                                                                                                                              |     3,170,304 |      12.68 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,512]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/transition2 (Linear)                                                                                 | weights: f32[48,512,128]         | f32[3,631,512]                                                                                                                                                                                                                                                                             | f32[3,631,128]                                                                                                                                                                                                                                              |     3,151,872 |      12.61 MB |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition/sequential (Sequential)                                                                           | bias: f32[48,128]                |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration/pair_transition (Transition)                                                                                      |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/evoformer_iteration (EvoformerIteration)                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer/__layer_stack_no_state_1 (_LayerStackNoState)                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/evoformer/single_activations (Linear)                                                                                                              | weights: f32[256,384]            | f32[631,256]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |        98,688 |     394.75 KB |
|  └ alphafold/alphafold_iteration/evoformer (EmbeddingsAndEvoformer)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                                 |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'act': f32[631,384],                                                                                                                                                                                                                                       |     2,018,540 |       8.07 MB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |  'final_atom_mask': f32[631,37],                                                                                                                                                                                                                            |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |  'final_atom_positions': f32[631,37,3]}                                                                                                                                                                                                                     |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/single_layer_norm (LayerNorm)                                                                                                     | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/initial_projection (Linear)                                                                                                       | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/pair_layer_norm (LayerNorm)                                                                                                       | offset: f32[128]                 | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,128]                                                                                                                                                                                                                                            |           256 |       1.02 KB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              | scale: f32[128]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f2e0f350>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f0dc0d90>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f0838e10>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23f01cde90>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23efbf1b50>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23ef4b89d0>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23eed4fed0>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                    |                                  | {'act': f32[631,384],                                                                                                                                                                                                                                                                      | ({'act': f32[631,384],                                                                                                                                                                                                                                      |     1,869,676 |       7.48 MB |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |  'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))}                                                                 |   'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))},                                |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  | aatype=s32[631], initial_act=f32[631,384], is_training=False, safe_key=<alphafold.model.prng.SafeKey object at 0x7f23ee6229d0>, sequence_mask=f32[631,1], static_feat_2d=f32[631,631,128], update_rigid=True                                                                               |  {'rigid': Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])),                                 |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |   'sc': {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                    |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])), |               |               |
|                                                                                                                                                                                  |                                  |                                                                                                                                                                                                                                                                                            |          'unnormalized_angles_sin_cos': f32[631,7,2]}})                                                                                                                                                                                                     |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                                | trainable_point_weights: f32[12] | inputs_1d=f32[631,384], inputs_2d=f32[631,631,128], mask=f32[631,1], rigid=Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])) | f32[631,384]                                                                                                                                                                                                                                                |     1,255,896 |       5.02 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,4], y=f32[631,12,4], z=f32[631,12,4])                                                                                                                                                                                                |        55,452 |     221.81 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection/point_projection (Linear)                                             | weights: f32[384,12,12]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,12]                                                                                                                                                                                                                                              |        55,440 |     221.76 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_point_projection (PointProjection)                                                  | bias: f32[12,12]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/q_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/k_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/attention_2d (Linear)                                                                    | weights: f32[128,12]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,12]                                                                                                                                                                                                                                             |         1,548 |       6.19 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[12]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_scalar_projection (Linear)                                                             | weights: f32[384,12,16]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,16]                                                                                                                                                                                                                                              |        73,728 |     294.91 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                     |                                  | f32[631,384], Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                                              | Vec3Array(x=f32[631,12,8], y=f32[631,12,8], z=f32[631,12,8])                                                                                                                                                                                                |       110,892 |     443.57 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection/point_projection (Linear)                                             | weights: f32[384,12,24]          | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,12,24]                                                                                                                                                                                                                                              |       110,880 |     443.52 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/v_point_projection (PointProjection)                                                  | bias: f32[12,24]                 |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention/output_projection (Linear)                                                               | weights: f32[2112,384]           | f32[631,2112]                                                                                                                                                                                                                                                                              | f32[631,384]                                                                                                                                                                                                                                                |       811,392 |       3.25 MB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/invariant_point_attention (InvariantPointAttention)                                                             | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/attention_layer_norm (LayerNorm)                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition (Linear)                                                                                                | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_1 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_2 (Linear)                                                                                              | weights: f32[384,384]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |       147,840 |     591.36 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | bias: f32[384]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/transition_layer_norm (LayerNorm)                                                                                  | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                             |                                  | f32[631,384]                                                                                                                                                                                                                                                                               | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631]))                                             |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid/rigid (Linear)                                                                                          | weights: f32[384,6]              | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,6]                                                                                                                                                                                                                                                  |         2,310 |       9.24 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/quat_rigid (QuatRigid)                                                                                          | bias: f32[6]                     |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                              |                                  | Rigid3Array(rotation=Rot3Array(xx=f32[631], xy=f32[631], xz=f32[631], yx=f32[631], yy=f32[631], yz=f32[631], zx=f32[631], zy=f32[631], zz=f32[631]), translation=Vec3Array(x=f32[631], y=f32[631], z=f32[631])), [f32[631,384], f32[631,384]], s32[631]                                    | {'angles_sin_cos': f32[631,7,2],                                                                                                                                                                                                                            |       166,414 |     665.66 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |  'atom_pos': Vec3Array(x=f32[631,14], y=f32[631,14], z=f32[631,14]),                                                                                                                                                                                        |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |  'frames': Rigid3Array(rotation=Rot3Array(xx=f32[631,8], xy=f32[631,8], xz=f32[631,8], yx=f32[631,8], yy=f32[631,8], yz=f32[631,8], zx=f32[631,8], zy=f32[631,8], zz=f32[631,8]), translation=Vec3Array(x=f32[631,8], y=f32[631,8], z=f32[631,8])),         |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |  'unnormalized_angles_sin_cos': f32[631,7,2]}                                                                                                                                                                                                               |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection (Linear)                                                                          | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/input_projection_1 (Linear)                                                                        | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2 (Linear)                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock1_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/resblock2_1 (Linear)                                                                               | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain/unnormalized_angles (Linear)                                                                       | weights: f32[128,14]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,14]                                                                                                                                                                                                                                                 |         1,806 |       7.22 KB |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration/rigid_sidechain (MultiRigidSidechain)                                                                           | bias: f32[14]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module/fold_iteration (FoldIteration)                                                                                                 |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration/structure_module (StructureModule)                                                                                                              |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/distogram_head (DistogramHead)                                                                                                                     |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'bin_edges': f32[63], 'logits': f32[631,631,64]}                                                                                                                                                                                                           |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/distogram_head/half_logits (Linear)                                                                                                                | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/distogram_head (DistogramHead)                                                                                                                  | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/masked_msa_head (MaskedMsaHead)                                                                                                                    |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[252,631,22]}                                                                                                                                                                                                                                 |         5,654 |      22.62 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/masked_msa_head/logits (Linear)                                                                                                                    | weights: f32[256,22]             | f32[252,631,256]                                                                                                                                                                                                                                                                           | f32[252,631,22]                                                                                                                                                                                                                                             |         5,654 |      22.62 KB |
|  └ alphafold/alphafold_iteration/masked_msa_head (MaskedMsaHead)                                                                                                                 | bias: f32[22]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                            |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[631,50]}                                                                                                                                                                                                                                     |        73,010 |     292.04 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/input_layer_norm (LayerNorm)                                                                                                   | offset: f32[384]                 | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,384]                                                                                                                                                                                                                                                |           768 |       3.07 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | scale: f32[384]                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/act_0 (Linear)                                                                                                                 | weights: f32[384,128]            | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        49,280 |     197.12 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/act_1 (Linear)                                                                                                                 | weights: f32[128,128]            | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,128]                                                                                                                                                                                                                                                |        16,512 |      66.05 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[128]                   |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_lddt_head/logits (Linear)                                                                                                                | weights: f32[128,50]             | f32[631,128]                                                                                                                                                                                                                                                                               | f32[631,50]                                                                                                                                                                                                                                                 |         6,450 |      25.80 KB |
|  └ alphafold/alphafold_iteration/predicted_lddt_head (PredictedLDDTHead)                                                                                                         | bias: f32[50]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/experimentally_resolved_head (ExperimentallyResolvedHead)                                                                                          |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'logits': f32[631,37]}                                                                                                                                                                                                                                     |        14,245 |      56.98 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/experimentally_resolved_head/logits (Linear)                                                                                                       | weights: f32[384,37]             | f32[631,384]                                                                                                                                                                                                                                                                               | f32[631,37]                                                                                                                                                                                                                                                 |        14,245 |      56.98 KB |
|  └ alphafold/alphafold_iteration/experimentally_resolved_head (ExperimentallyResolvedHead)                                                                                       | bias: f32[37]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_aligned_error_head (PredictedAlignedErrorHead)                                                                                           |                                  | {'msa': f32[252,631,256],                                                                                                                                                                                                                                                                  | {'breaks': f32[63], 'logits': f32[631,631,64]}                                                                                                                                                                                                              |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |  'msa_first_row': f32[631,256],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |  'pair': f32[631,631,128],                                                                                                                                                                                                                                                                 |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'single': f32[631,384],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'structure_module': f32[631,384]}, {'aatype': s32[631],                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_mask': s32[631,37],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'all_atom_positions': f32[631,37,3],                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'assembly_num_chains': s32[],                                                                                                                                                                                                                                                             |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'asym_id': f32[631],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'bert_mask': f32[6143,631],                                                                                                                                                                                                                                                               |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'cluster_bias_mask': f32[6143],                                                                                                                                                                                                                                                           |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_matrix': f32[6143,631],                                                                                                                                                                                                                                                         |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'deletion_mean': f32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_id': f32[631],                                                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'entity_mask': s32[631],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa': s32[6143,631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'msa_mask': f32[6143,631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_alignments': s32[],                                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'num_templates': s32[],                                                                                                                                                                                                                                                                   |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_msa_first_row': f32[631,256],                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pair': f32[631,631,128],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'prev_pos': f32[631,37,3],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'residue_index': s32[631],                                                                                                                                                                                                                                                                |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_length': s32[],                                                                                                                                                                                                                                                                      |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'seq_mask': f32[631],                                                                                                                                                                                                                                                                     |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'sym_id': f32[631],                                                                                                                                                                                                                                                                       |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_aatype': s32[4,631],                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_mask': f32[4,631,37],                                                                                                                                                                                                                                                  |                                                                                                                                                                                                                                                             |               |               |
|                                                                                                                                                                                  |                                  |  'template_all_atom_positions': f32[4,631,37,3]}, False                                                                                                                                                                                                                                    |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| alphafold/alphafold_iteration/predicted_aligned_error_head/logits (Linear)                                                                                                       | weights: f32[128,64]             | f32[631,631,128]                                                                                                                                                                                                                                                                           | f32[631,631,64]                                                                                                                                                                                                                                             |         8,256 |      33.02 KB |
|  └ alphafold/alphafold_iteration/predicted_aligned_error_head (PredictedAlignedErrorHead)                                                                                        | bias: f32[64]                    |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold/alphafold_iteration (AlphaFoldIteration)                                                                                                                            |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
|  └ alphafold (AlphaFold)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+