+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+----------------------------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+---------------+---------------+
| 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)                                                                                                                                                         |                                  |                                                                                                                                                                                                                                                                                            |                                                                                                                                                                                                                                                             |               |               |
+---------